#ifndef PRECOND_HPP
#define PRECOND_HPP

#include <stdio.h>
#include "omp.h"
#include "my_parcsr.hpp"

/// @brief 预条件子的基类
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 数据类型，可以为float或者double
/// @tparam calc_t 计算类型，可以为float或者double
/// @tparam setup_t 建立类型，可以为float或者double
template<typename idx_t, typename data_t, typename calc_t, typename setup_t>
class Precond
{
public:
    /// @brief 预条件子的构造函数
    Precond() {}
    /// @brief 预条件子的析构函数
    virtual ~Precond() {};
    /// @brief 预条件子的建立函数，需要在使用之前调用
    /// @param A 系数矩阵
    virtual void Setup(const MyParCSRMatrix<idx_t, setup_t, setup_t> & A) = 0;
    /// @brief 预条件子的求解函数
    /// @param b 右端项
    /// @param x 解向量
    virtual void Solve(const MyParCSRVector<idx_t, calc_t> & b, MyParCSRVector<idx_t, calc_t> & x) = 0;
};

/// @brief Jacobi预条件子
template<typename idx_t, typename data_t, typename calc_t, typename setup_t>
class Jacobi: public Precond<idx_t, data_t, calc_t, setup_t>
{
private:
    idx_t n; // 矩阵的局部行数
    data_t *diag = nullptr; // 对角线元素
public:
    /// @brief Jacobi预条件子的构造函数
    Jacobi() {}
    /// @brief Jacobi预条件子的析构函数
    /// @note 释放内存
    ~Jacobi() { if (diag) free(diag); }
    /// @brief Jacobi预条件子的建立函数
    /// @param A 系数矩阵
    /// @note 仅仅需要提取对角线元素
    virtual void Setup(const MyParCSRMatrix<idx_t, setup_t, setup_t> & A) {
        this->n = A.end_row - A.beg_row;
        if (diag) free(diag);
        diag = (data_t *) malloc(sizeof(data_t) * n);

        // 获取 A 的对角线块
        MyCSRMatrix<idx_t, setup_t, setup_t> * A_diag = A.diag;
        idx_t * row_ptr = A_diag->row_ptr;
        idx_t * col_idx = A_diag->col_idx;
        setup_t * vals = A_diag->vals;

        // 提取对角线元素
        #pragma omp parallel for schedule(static)
        for (idx_t i = 0; i < n; i++) {
            idx_t begin = row_ptr[i], end = row_ptr[i + 1];
            for (int j = begin; j < end; j++) {
                if (i == col_idx[j]) {
                    diag[i] = (data_t) (1.0 / vals[j]);
                    break;
                }
            }
        }
    }
    /// @brief Jacobi预条件子的求解函数
    /// @param b 右端项
    /// @param x 解向量
    /// @note x = D^{-1} b
    virtual void Solve(const MyParCSRVector<idx_t, calc_t> & b, MyParCSRVector<idx_t, calc_t> & x) {
        const idx_t nrows = b.end_row - b.beg_row;
        assert(nrows <= this->n);
        #pragma omp parallel for schedule(static)
        for (idx_t i = 0; i < nrows; i++) {
            x.data[i] = b.data[i] * diag[i];
        }
    }
};

/// @brief Point Gauss-Seidel预条件子
template<typename idx_t, typename data_t, typename calc_t, typename setup_t>
class PGS: public Precond<idx_t, data_t, calc_t, setup_t>
{
private:
    idx_t n = 0; // 矩阵的局部行数
    bool inited = false; // 是否已经初始化

    // 下三角 CSR
    idx_t *row_ptr_lower = nullptr;
    idx_t *col_idx_lower = nullptr;
    data_t *val_lower = nullptr;

    // 上三角 CSR
    idx_t *row_ptr_upper = nullptr;
    idx_t *col_idx_upper = nullptr;
    data_t *val_upper = nullptr;

    // 对角线元素
    data_t *val_diag = nullptr;
public:
    /// @brief PGS预条件子的构造函数
    PGS() {}
    /// @brief PGS预条件子的析构函数
    /// @note 释放内存
    ~PGS() {
        if (inited) {
            free(row_ptr_lower);
            free(col_idx_lower);
            free(val_lower);
            free(row_ptr_upper);
            free(col_idx_upper);
            free(val_upper);
            free(val_diag);
        }
    }
    /// @brief PGS预条件子的建立函数
    /// @param A 系数矩阵
    /// @note 提取并分离 下三角、上三角、对角线元素
    virtual void Setup(const MyParCSRMatrix<idx_t, setup_t, setup_t> & A) {
        this->n = A.end_row - A.beg_row;
        if (inited) {
            free(row_ptr_lower);
            free(col_idx_lower);
            free(val_lower);
            free(row_ptr_upper);
            free(col_idx_upper);
            free(val_upper);
            free(val_diag);
        }

        // 获取 A 的对角线块
        MyCSRMatrix<idx_t, setup_t, setup_t> * A_diag = A.diag;
        idx_t * row_ptr = A_diag->row_ptr;
        idx_t * col_idx = A_diag->col_idx;
        setup_t * vals = A_diag->vals;

        // 分配内存
        row_ptr_lower = (idx_t *) malloc(sizeof(idx_t) * (n + 1));
        row_ptr_upper = (idx_t *) malloc(sizeof(idx_t) * (n + 1));

        // 初始化 row_ptr，不可省略！
        #pragma omp parallel for schedule(static)
        for (int i = 0; i <= n; i++)
            row_ptr_lower[i] = row_ptr_upper[i] = 0;

        // 统计下三角、上三角元素个数
        #pragma omp parallel for schedule(static)
        for (idx_t i = 0; i < n; i++) {
            idx_t begin = row_ptr[i], end = row_ptr[i + 1];
            for (idx_t j = begin; j < end; j++) {
                idx_t col = col_idx[j];
                if (col < i) {
                    row_ptr_lower[i + 1]++;
                } else if (i < col) {
                    row_ptr_upper[i + 1]++;
                }
            }
        }

        // 计算下三角、上三角的行偏移量
        for (idx_t i = 0; i < n; i++) {
            row_ptr_lower[i+1] += row_ptr_lower[i];
            row_ptr_upper[i+1] += row_ptr_upper[i];
        }

        // 分配内存
        idx_t nnz_lower = row_ptr_lower[n];
        idx_t nnz_upper = row_ptr_upper[n];
        col_idx_lower = (idx_t *) malloc(sizeof(idx_t) * nnz_lower);
        val_lower = (data_t *) malloc(sizeof(data_t) * nnz_lower);
        col_idx_upper = (idx_t *) malloc(sizeof(idx_t) * nnz_upper);
        val_upper = (data_t *) malloc(sizeof(data_t) * nnz_upper);
        val_diag = (data_t *) malloc(sizeof(data_t) * n);

        // 提取下三角、上三角、对角线元素
        #pragma omp parallel for schedule(static)
        for (idx_t i = 0; i < n; i++) {
            idx_t begin = row_ptr[i], end = row_ptr[i + 1];
            idx_t ptr_lower = row_ptr_lower[i];
            idx_t ptr_upper = row_ptr_upper[i];
            for (idx_t j = begin; j < end; j++) {
                idx_t col = col_idx[j];
                setup_t v = vals[j];
                if (col == i) { // 对角线元素
                    val_diag[i] = (data_t) (1.0 / v); 
                } else if (col < i) { // 下三角元素
                    col_idx_lower[ptr_lower] = col;
                    val_lower[ptr_lower] = (data_t) v;
                    ptr_lower++;
                } else if (i < col){ // 上三角元素
                    col_idx_upper[ptr_upper] = col;
                    val_upper[ptr_upper] = (data_t) v;
                    ptr_upper++;
                }
            }   
        }
        inited = true;
    }
    /// @brief PGS预条件子的求解函数
    /// @param b 右端项
    /// @param x 解向量
    /// @note 使用了零初值优化
    /// @note 详见文档
    virtual void Solve(const MyParCSRVector<idx_t, calc_t> & b, MyParCSRVector<idx_t, calc_t> & x) {
        const idx_t nrows = b.end_row - b.beg_row;
        assert(nrows <= this->n);
        calc_t res = 0.0, zero = 0.0;
        calc_t * x_data = x.data;
        const calc_t * b_data = b.data;

        // 前扫：D*x^{t+1} = b - L*x^{t+1}
        for (idx_t i = 0; i < nrows; i++) {
            res = b_data[i];
            idx_t begin = row_ptr_lower[i], end = row_ptr_lower[i + 1];
            for (idx_t j = begin; j < end; j++) {
                res -= ((calc_t)val_lower[j]) * x_data[col_idx_lower[j]];
            }
            x_data[i] = res * val_diag[i];
        }

        // 后扫：D*x^{t+2} = D*x^{t+1} - U*x^{t+2}
        for (idx_t i = nrows - 1; i >= 0; i--) {
            res = 0.0;
            idx_t begin = row_ptr_upper[i], end = row_ptr_upper[i + 1];
            for (idx_t j = begin; j < end; j++) {
                res += ((calc_t)val_upper[j]) * x_data[col_idx_upper[j]];
            }
            x_data[i] -= res * val_diag[i];
        }
    }
};

#endif // PRECOND_HPP